import importlib
import os
import logging
import sys
import threading

import grpc
from signal import signal, SIGTERM
from concurrent import futures
from service.pb.model_server_pb2 import DatasetDownloadRequest
from service.pb.model_server_pb2_grpc import ModelServerServicer, add_ModelServerServicer_to_server
from loguru import logger
from service.tools import get_file_chunk, debug_exit


class ModelServer(ModelServerServicer):
    def __init__(self):
        self.project_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        os.chdir(self.project_path)


    @debug_exit
    def dataset_download(self, request: DatasetDownloadRequest, context):
        dataset_name = request.dataset_name
        mode = request.mode
        task = request.task
        assert dataset_name and mode and task
        logger.info("{} Dataset {} {} download request for ".format(task, dataset_name, mode))
        dataset_module = importlib.import_module('database.{}.{}.interface'.format(task, dataset_name))
        full_file_name = dataset_module.get_file_name(mode)
        return get_file_chunk(full_file_name)


def serve():
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    add_ModelServerServicer_to_server(
        ModelServer(), server)
    server.add_insecure_port('[::]:50051')
    server.start()
    logger.info("Server start at 50051")

    def handle_sigterm(signum, frame):
        logger.info('Got signal {}, {}'.format(signum, frame))
        server.stop(None)
        logger.critical("Stop Server")

    signal(SIGTERM, handle_sigterm)
    server.wait_for_termination()


if __name__ == '__main__':
    logging.basicConfig()
    serve()
